@echo off
REM FLUX 24G显存微调训练脚本 - Windows版本
REM 适用于24GB VRAM的GPU

echo ========================================
echo FLUX 24G显存微调训练脚本 - Windows版本
echo ========================================

REM 设置环境变量
set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
set CUDA_LAUNCH_BLOCKING=1
set TOKENIZERS_PARALLELISM=false

REM 基础路径定义 - 适配当前环境
set WORKSPACE=E:\AI
set AI_DIR=%WORKSPACE%
set PYTHON_PATH=%AI_DIR%\python\python.exe
set SCRIPT_PATH=%AI_DIR%\ComfyUI\custom_nodes\comfyui_lora_train\sd_scripts\flux_train.py

REM 模型路径定义
set MODELS_DIR=%AI_DIR%\ComfyUI\models
set FLUX_MODEL=%MODELS_DIR%\checkpoints\flux1-dev.safetensors
set CLIP_L_MODEL=%MODELS_DIR%\clip\clip_l.safetensors
set T5XXL_MODEL=%MODELS_DIR%\clip\t5xxl_fp8_e4m3fn.safetensors
set AE_MODEL=%MODELS_DIR%\vae\ae.safetensors

REM 数据集和输出路径
set TRAIN_DATA_DIR=%AI_DIR%\ComfyUI\custom_nodes\comfyui_lora_train\training_data_all\images
set OUTPUT_DIR=%AI_DIR%\ComfyUI\custom_nodes\comfyui_lora_train\training_data_all\output
set SAMPLE_PROMPTS=%AI_DIR%\ComfyUI\custom_nodes\comfyui_lora_train\training_data_all\sample_prompts.txt

echo 检查必要文件...

if not exist "%FLUX_MODEL%" (
    echo 错误: FLUX模型文件不存在: %FLUX_MODEL%
    echo 请确保flux1-dev.safetensors文件在%MODELS_DIR%\checkpoints\目录
    pause
    exit /b 1
)

if not exist "%CLIP_L_MODEL%" (
    echo 错误: CLIP-L模型文件不存在: %CLIP_L_MODEL%
    echo 请确保clip_l.safetensors文件在%MODELS_DIR%\clip\目录
    pause
    exit /b 1
)

if not exist "%T5XXL_MODEL%" (
    echo 错误: T5XXL模型文件不存在: %T5XXL_MODEL%
    echo 请确保t5xxl_fp8_e4m3fn.safetensors文件在%MODELS_DIR%\clip\目录
    pause
    exit /b 1
)

if not exist "%AE_MODEL%" (
    echo 错误: AE模型文件不存在: %AE_MODEL%
    echo 请确保ae.safetensors文件在%MODELS_DIR%\vae\目录
    pause
    exit /b 1
)

if not exist "%TRAIN_DATA_DIR%" (
    echo 错误: 训练数据目录不存在: %TRAIN_DATA_DIR%
    echo 请确保training_data_all\images目录存在并包含训练图像
    pause
    exit /b 1
)

if not exist "%SAMPLE_PROMPTS%" (
    echo 错误: 采样提示词文件不存在: %SAMPLE_PROMPTS%
    echo 请确保sample_prompts.txt文件在training_data_all\目录
    pause
    exit /b 1
)

REM 检查Python环境
echo 检查Python环境...
if not exist "%PYTHON_PATH%" (
    echo 错误: Python路径不存在: %PYTHON_PATH%
    pause
    exit /b 1
)

REM 检查accelerate
echo 检查accelerate...
"%PYTHON_PATH%" -c "import accelerate" 2>nul
if %ERRORLEVEL% neq 0 (
    echo 错误: accelerate未安装
    echo 请运行: "%PYTHON_PATH%" -m pip install accelerate
    pause
    exit /b 1
)

REM 检查GPU
echo 检查GPU...
nvidia-smi >nul 2>&1
if %ERRORLEVEL% neq 0 (
    echo 错误: nvidia-smi未找到，请确保NVIDIA驱动已安装
    pause
    exit /b 1
)

REM 检查GPU显存
echo 检查GPU显存...
for /f "tokens=*" %%i in ('nvidia-smi --query-gpu^=memory.total --format^=csv,noheader,nounits') do set GPU_MEMORY=%%i
if %GPU_MEMORY% LSS 20000 (
    echo 警告: GPU显存不足24GB (当前: %GPU_MEMORY%MB)
    echo 建议使用24GB或更大显存的GPU
    set /p CONTINUE="是否继续训练? (y/N): "
    if /i not "%CONTINUE%"=="y" (
        echo 训练已取消
        pause
        exit /b 1
    )
)

REM 创建输出目录
if not exist "%OUTPUT_DIR%" (
    echo 创建输出目录: %OUTPUT_DIR%
    mkdir "%OUTPUT_DIR%"
)

echo ========================================
echo 训练配置信息
echo ========================================
echo 模型: %FLUX_MODEL%
echo 训练数据: %TRAIN_DATA_DIR%
echo 输出目录: %OUTPUT_DIR%
echo 采样提示词: %SAMPLE_PROMPTS%
echo GPU显存: %GPU_MEMORY%MB
echo.

REM 显示训练前GPU显存使用情况
echo 训练前GPU显存使用情况：
nvidia-smi --query-gpu=memory.total,memory.used,memory.free --format=csv,noheader,nounits

echo ========================================
echo 开始FLUX微调训练
echo ========================================

REM 设置环境变量确保使用正确的Python
set PATH=%AI_DIR%\python;%PATH%

REM 微调训练命令
accelerate launch --mixed_precision bf16 --num_cpu_threads_per_process 1 --num_processes 1 --num_machines 1 --dynamo_backend no "%SCRIPT_PATH%" ^
--pretrained_model_name_or_path "%FLUX_MODEL%" ^
--clip_l "%CLIP_L_MODEL%" ^
--t5xxl "%T5XXL_MODEL%" ^
--ae "%AE_MODEL%" ^
--train_data_dir "%TRAIN_DATA_DIR%" ^
--output_dir "%OUTPUT_DIR%" ^
--output_name flux-finetune-24g ^
--learning_rate 5e-5 ^
--max_train_epochs 4 ^
--train_batch_size 1 ^
--resolution 1024 ^
--optimizer_type adafactor ^
--optimizer_args "relative_step=False" "scale_parameter=False" "warmup_init=False" ^
--lr_scheduler constant_with_warmup ^
--mixed_precision bf16 ^
--save_precision bf16 ^
--save_every_n_epochs 1 ^
--save_model_as safetensors ^
--gradient_checkpointing ^
--max_data_loader_n_workers 2 ^
--persistent_data_loader_workers ^
--seed 42 ^
--max_grad_norm 0.0 ^
--timestep_sampling shift ^
--discrete_flow_shift 3.1582 ^
--model_prediction_type raw ^
--guidance_scale 1.0 ^
--fused_backward_pass ^
--blocks_to_swap 8 ^
--disable_mmap_load_safetensors ^
--cache_latents_to_disk ^
--cache_text_encoder_outputs_to_disk ^
--sample_every_n_steps 1000 ^
--sample_prompts "%SAMPLE_PROMPTS%" ^
--sample_sampler ddim ^
--caption_extension .txt

if %ERRORLEVEL% neq 0 (
    echo 错误: 训练失败，退出代码: %ERRORLEVEL%
    pause
    exit /b %ERRORLEVEL%
)

echo ========================================
echo FLUX微调训练完成！
echo ========================================

REM 显示最终显存使用情况
echo 训练后GPU显存使用情况：
nvidia-smi --query-gpu=memory.total,memory.used,memory.free --format=csv,noheader,nounits

REM 检查输出文件
echo 检查输出文件...
if exist "%OUTPUT_DIR%\flux-finetune-24g" (
    echo ✓ 模型输出目录已创建: %OUTPUT_DIR%\flux-finetune-24g
    dir "%OUTPUT_DIR%\flux-finetune-24g"
    
    REM 检查样本图像
    if exist "%OUTPUT_DIR%\flux-finetune-24g\sample_images" (
        echo ✓ 样本图像目录已创建
        for /f %%i in ('dir /b "%OUTPUT_DIR%\flux-finetune-24g\sample_images\*.png" 2^>nul ^| find /c /v ""') do set SAMPLE_COUNT=%%i
        echo 样本图像数量: %SAMPLE_COUNT%
    )
    
    REM 检查训练日志
    if exist "%OUTPUT_DIR%\flux-finetune-24g\train_log.txt" (
        echo ✓ 训练日志已生成
        for %%A in ("%OUTPUT_DIR%\flux-finetune-24g\train_log.txt") do echo 日志文件大小: %%~zA bytes
    )
) else (
    echo 警告: 模型输出目录未找到
)

echo 训练完成！模型文件保存在: %OUTPUT_DIR%\flux-finetune-24g\
echo 您可以在ComfyUI中加载生成的微调模型进行测试
pause
